import os

# os.environ['USE_PATH_FOR_GDAL_PYTHON'] = 'YES'
os.environ["PATH"] = "D:\\develop\\Anaconda3\\envs\\geo_env\\Lib\\site-packages\\osgeo;" + os.environ["PATH"]
os.environ["PYTHONPATH"] = "D:\\develop\\Anaconda3\\envs\\geo_env\\Lib\\site-packages\\osgeo;" + os.environ["PATH"]
from osgeo import gdal, ogr, osr
import numpy as np


def toPolygonize(path):
    inraster = gdal.Open(path)  # 读取路径中的栅格数据
    inband = inraster.GetRasterBand(1)  # 这个波段就是最后想要转为矢量的波段，如果是单波段数据的话那就都是1
    #
    # 设置掩膜值
    inband.SetNoDataValue(0.0)
    # inband.SetNoDataValue(np.NaN)

    # 读取波段数组
    # inbandData=inband.ReadAsArray()
    # inbandData[inbandData < 100] = 11
    # inband.WriteArray(inbandData)
    # 读取波段掩膜数组
    # mask_band_array = inband.GetMaskBand().ReadAsArray()

    # inband.CreateMaskBand(0)
    prj = osr.SpatialReference()
    prj.ImportFromWkt(inraster.GetProjection())  # 读取栅格数据的投影信息，用来为后面生成的矢量做准备
    #
    outshp = path[:-4] + "5.shp"
    drv = ogr.GetDriverByName("ESRI Shapefile")
    if os.path.exists(outshp):  # 若文件已经存在，则删除它继续重新做一遍
        drv.DeleteDataSource(outshp)
    Polygon = drv.CreateDataSource(outshp)  # 创建一个目标文件
    Poly_layer = Polygon.CreateLayer(path, srs=prj, geom_type=ogr.wkbMultiPolygon)
    newField = ogr.FieldDefn('value',
                             ogr.OFTInteger)  # 给目标shp文件添加一个字段，用来存储原始栅格的pixel value,OFTReal浮点型，ogr.OFTInteger是整数型
    Poly_layer.CreateField(newField)
    # inband.SetNoDataValue(-1000)
    gdal.Polygonize(inband, inband.GetMaskBand(), Poly_layer, 0)  # 核心函数，执行的就是栅格转矢量操作
    # gdal.FPolygonize(inband, None, Poly_layer, 0)  # 只转矩形，不合并
    Polygon.SyncToDisk()
    Polygon = None


toPolygonize('G:\\code\\ydh\\py_tiff_back\\testData\\tiffData\\RS_FOR_BJ_FMP_202409030610_202409031110_1H_5_1_1.tif')
